#ifndef CONV2D_H
#define CONV2D_H

#include "type_alias.h"
#include <cstring>
#include <memory>

namespace wawa {
    class Conv2d {
    public:
        Conv2d(size_t in_channels, size_t out_channels,
               size_t kernel_size, size_t stride = 1,
               size_t padding = 0, bool bias = false);


        Tensor m_weights;
        size_t m_in_channels;
        size_t m_out_channels;
        size_t m_kernel_size;
        size_t m_stride;
        size_t m_padding;
        bool m_bias;
        std::shared_ptr<double[]> m_bias_weights;


    public:
        [[nodiscard]] std::pair<size_t, size_t> calculate_output_size(size_t H_in, size_t W_in) const;

        [[nodiscard]] Tensor forward(const Tensor &input) const;

        void initialize_parameters() const;

        void convolve(const Tensor &input, Tensor &output) const;
    };
}
#endif //!CONV2D_H
